#!/usr/bin/env python3
"""
AI对话功能修复最终验证脚本
验证服务器是否只返回text字段内容，而不是整个JSON对象
"""

import asyncio
import aiohttp
import json
import sys

async def test_ai_chat_stream():
    """测试流式AI对话响应格式"""
    print("=== 测试流式AI对话响应格式 ===")
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                'http://127.0.0.1:8003/ai/chat',
                json={
                    "prompt": "你好，请简单介绍一下硅材料",
                    "temperature": 0.7,
                    "max_tokens": 100
                }
            ) as response:
                
                if response.status != 200:
                    print(f"❌ 请求失败，状态码: {response.status}")
                    return False
                
                print(f"✅ 请求成功，状态码: {response.status}")
                
                # 读取流式响应
                full_content = ""
                async for line in response.content:
                    line_text = line.decode('utf-8').strip()
                    if line_text and line_text.startswith('data: '):
                        json_str = line_text[6:]  # 移除 'data: ' 前缀
                        
                        try:
                            data = json.loads(json_str)
                            print(f"\n收到数据块: {json.dumps(data, ensure_ascii=False)}")
                            
                            # 验证数据格式
                            keys = list(data.keys())
                            has_content = 'content' in data
                            has_is_final = 'is_final' in data
                            has_only_expected_fields = len(keys) <= 2 and has_content
                            
                            if has_only_expected_fields:
                                print("✅ 数据格式正确：只包含content和is_final字段")
                                if data['content']:
                                    full_content += data['content']
                                    print(f"✅ 成功提取text内容: {data['content']}")
                            else:
                                print(f"❌ 数据格式错误：包含额外字段 {keys}")
                                return False
                                
                            if data.get('is_final'):
                                print("✅ 收到结束标记")
                                break
                                
                        except json.JSONDecodeError as e:
                            print(f"❌ JSON解析错误: {e}")
                            return False
                
                print(f"\n✅ 完整响应内容: {full_content}")
                return True
                
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False

async def test_ai_chat_simple():
    """测试非流式AI对话响应格式"""
    print("\n=== 测试非流式AI对话响应格式 ===")
    
    try:
        async with aiohttp.ClientSession() as session:
            async with session.post(
                'http://127.0.0.1:8003/ai/chat/simple',
                json={
                    "prompt": "你好，请简单介绍一下碳材料",
                    "temperature": 0.7,
                    "max_tokens": 100
                }
            ) as response:
                
                if response.status != 200:
                    print(f"❌ 请求失败，状态码: {response.status}")
                    return False
                
                data = await response.json()
                print(f"✅ 请求成功，响应数据: {json.dumps(data, ensure_ascii=False, indent=2)}")
                
                # 验证数据格式
                keys = list(data.keys())
                has_response = 'response' in data
                has_only_response = len(keys) == 1 and has_response
                
                if has_only_response:
                    print("✅ 数据格式正确：只包含response字段")
                    if data['response']:
                        print(f"✅ 成功提取text内容: {data['response']}")
                    return True
                else:
                    print(f"❌ 数据格式错误：包含额外字段 {keys}")
                    return False
                
    except Exception as e:
        print(f"❌ 测试失败: {e}")
        return False

async def simulate_original_problem():
    """模拟原始问题：显示整个JSON对象"""
    print("\n=== 模拟原始问题验证修复 ===")
    
    # 模拟用户之前看到的完整JSON响应
    original_problem_response = {
        "id": "cmpl-f4c15c81fb1b4740a91ad40902d3ee53",
        "object": "text_completion", 
        "created": 1758753075,
        "model": "S1-Base-32B",
        "choices": [
            {
                "index": 0,
                "text": "长依",
                "logprobs": None,
                "finish_reason": None,
                "stop_reason": None,
                "prompt_token_ids": None,
                "token_ids": None
            }
        ],
        "usage": None
    }
    
    print(f"原始问题示例（用户之前看到的）:")
    print(json.dumps(original_problem_response, ensure_ascii=False, indent=2))
    
    # 验证修复后的响应格式
    print(f"\n修复后的响应格式（现在应该看到的）:")
    fixed_response = {
        "content": "长依",
        "is_final": True
    }
    print(json.dumps(fixed_response, ensure_ascii=False, indent=2))
    
    # 对比验证
    original_keys = set(original_problem_response.keys())
    fixed_keys = set(fixed_response.keys())
    
    print(f"\n对比验证:")
    print(f"原始响应字段数: {len(original_keys)}")
    print(f"修复后响应字段数: {len(fixed_keys)}")
    print(f"字段减少: {len(original_keys) - len(fixed_keys)} 个")
    
    if fixed_keys.issubset({'content', 'is_final'}):
        print("✅ 修复成功：响应只包含content和is_final字段")
        return True
    else:
        print("❌ 修复失败：响应包含额外字段")
        return False

async def main():
    """主测试函数"""
    print("AI对话功能修复验证")
    print("=" * 50)
    
    # 检查服务器连接
    try:
        async with aiohttp.ClientSession() as session:
            async with session.get('http://127.0.0.1:8003/health') as response:
                if response.status == 200:
                    print("[OK] 服务器连接正常")
                else:
                    print("[ERROR] 服务器连接失败")
                    return
    except Exception as e:
        print(f"[ERROR] 无法连接到服务器: {e}")
        return
    
    # 运行测试
    tests_passed = 0
    total_tests = 3
    
    # 测试1: 流式响应格式
    if await test_ai_chat_stream():
        tests_passed += 1
    
    # 测试2: 非流式响应格式
    if await test_ai_chat_simple():
        tests_passed += 1
        
    # 测试3: 模拟原始问题验证
    if await simulate_original_problem():
        tests_passed += 1
    
    # 总结
    print(f"\n" + "=" * 50)
    print(f"测试结果: {tests_passed}/{total_tests} 通过")
    
    if tests_passed == total_tests:
        print("[SUCCESS] AI对话功能修复验证成功！")
        print("[OK] 服务器现在只返回text字段内容，而不是整个JSON对象")
        return 0
    else:
        print("[ERROR] AI对话功能修复验证失败")
        return 1

if __name__ == "__main__":
    result = asyncio.run(main())
    sys.exit(result)